#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys

sys.path.append('.')

"""
   对select merge的人脸ID进行替换
"""


def read_file_list(file_list):
    faceid_subfaceid_files = []
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            # file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)

            faceid_subfaceid_files.append([faceid_subfaceid, line])

        print('Read File Count: {}'.format(file_count))
    return faceid_subfaceid_files


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            # file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{} {}'.format(faceid, subfaceid)
            feature = feat_file.read(feat_dim*4)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                faceid_subfaceid_file_features[faceid_subfaceid] = [[line], [feature]]
            else:
                file_features[0].append(line)
                file_features[1].append(feature)

        print('Read File Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def read_merge_file(file_list):
    faceid_subfaceids_map = {}
    delete_faceid_subfaceids= set()
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # NONE:960362bc-3dec-4785-8c87-bd60052bb7f0_192a5bc2-ef6f-4ef5-b82f-9bc6cd0791eb,960362bc-3dec-4785-8c87-bd60052bb7f0_29685424-4d56-49fb-acb7-3416fba1bb15,960362bc-3dec-4785-8c87-bd60052bb7f0_dc9d79c5-cccf-4e21-8ece-e2a14f00325e
            lines = line.strip().split(':')
            faceid_nums = lines[1].split(';')
            faceids = faceid_nums[0].split(',')
            faceids = [faceid.replace('_', ' ', 1) for faceid in faceids]

            facenums = faceid_nums[1].split(',')
            facenums = list(map(int, facenums))

            num_faceids = {}
            for i, num in enumerate(facenums):
                if num == 0:
                    continue
                elif num < 0:
                    delete_faceid_subfaceids.add(faceids[i])
                else:
                    _faceids = num_faceids.get(num, None)
                    if _faceids is None:
                        num_faceids[num] = [faceids[i]]
                    else:
                        _faceids.append(faceids[i])

            for num, faceids in num_faceids.items():
                faceid = faceids[0]
                for replace_faceid in faceids[1:]:
                    faceid_subfaceids_map[replace_faceid] = faceid
        print('Read Merge Count: {}'.format(file_count))
    return faceid_subfaceids_map, delete_faceid_subfaceids


def main():
    parser = argparse.ArgumentParser(description="Replace Face ReID Image")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/FilterSelectFace.txt",
        help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/FilterSelectFeature.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--merge_file", default="/rootfs/media/kasim/Data/data/FaceImage/filter/select_merge_faces.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage/filter", help="path to output", type=str
    )
    parser.add_argument("--feature_dim", default=256, help="feature dim", type=int)
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim)
    faceid_subfaceids_map, delete_faceid_subfaceids = read_merge_file(args.merge_file)

    out_file_path = os.path.join(output_dir, 'MR' + os.path.basename(args.file_list))
    out_feat_file_path = os.path.join(output_dir, 'MR' + os.path.basename(args.feature_file))
    out_file = open(out_file_path, 'w')
    out_feat_file = open(out_feat_file_path, 'wb')
    file_count = 0
    for faceid_subfaceid, file_features in faceid_subfaceid_file_features.items():
        if faceid_subfaceid in delete_faceid_subfaceids:
            continue
        replace_faceid_subfaceid = faceid_subfaceids_map.get(faceid_subfaceid, None)
        files, features = file_features
        if replace_faceid_subfaceid is None:
            for file, feature in zip(files, features):
                out_file.write(file)
                out_feat_file.write(feature)
        else:
            for file, feature in zip(files, features):
                file = file.replace(faceid_subfaceid, replace_faceid_subfaceid, 1)
                out_file.write(file)
                out_feat_file.write(feature)
        file_count += 1
        if file_count % 1000 == 0:
            print('Write Count: {}'.format(file_count))
    print('Write Count: {}'.format(file_count))
    out_file.close()
    out_feat_file.close()
    os.system('chmod a+wr {}'.format(out_file_path))
    os.system('chmod a+wr {}'.format(out_feat_file_path))
    print('finish!')


if __name__ == '__main__':
    main()
